import cv2
import os
import numpy as np


def slice_image_with_overlap(image_path, output_dir, window_size=640, step_size=80):
    # 读取大图片
    image = cv2.imread(image_path)
    h, w, _ = image.shape  # 获取图片的高度和宽度

    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 初始化切片计数器
    count = 0

    # 使用滑动窗口遍历图片
    for y in range(0, h - window_size + 1, step_size):
        for x in range(0, w - window_size + 1, step_size):
            # 计算切片的窗口范围，确保切片是 640x640
            slice_img = image[y:y + window_size, x:x + window_size]

            # 保存切片
            slice_filename = os.path.join(output_dir, f"slice_{count}.png")
            cv2.imwrite(slice_filename, slice_img)
            print(f"Saved {slice_filename}")
            count += 1

    # 处理边缘部分：确保最后的一些区域也被覆盖到
    # 如果宽度不能被整除，就在宽度方向上补充切片
    if w % step_size != 0:
        for y in range(0, h - window_size + 1, step_size):
            slice_img = image[y:y + window_size, w - window_size:w]
            slice_filename = os.path.join(output_dir, f"slice_{count}.png")
            cv2.imwrite(slice_filename, slice_img)
            print(f"Saved {slice_filename}")
            count += 1

    # 如果高度不能被整除，就在高度方向上补充切片
    if h % step_size != 0:
        for x in range(0, w - window_size + 1, step_size):
            slice_img = image[h - window_size:h, x:x + window_size]
            slice_filename = os.path.join(output_dir, f"slice_{count}.png")
            cv2.imwrite(slice_filename, slice_img)
            print(f"Saved {slice_filename}")
            count += 1

    # 处理右下角的剩余区域
    if w % step_size != 0 and h % step_size != 0:
        slice_img = image[h - window_size:h, w - window_size:w]
        slice_filename = os.path.join(output_dir, f"slice_{count}.png")
        cv2.imwrite(slice_filename, slice_img)
        print(f"Saved {slice_filename}")


# 使用函数进行切片
image_path = 'black_template.jpg'  # 输入大图片的路径
output_dir = 'rawB'  # 输出切片的保存路径

slice_image_with_overlap(image_path, output_dir, window_size=640, step_size=160)
